高阶优化器：深度学习加速的利器

摘要

深度学习的训练往往极度依赖优化器的选择，从而获取最佳的模型参数来进行相关任务的预测。随着深度学习的高速发展，尤其是一阶优化器SGD的成功，使得越来越多的研究者对高阶优化器产生了兴趣。本文将介绍最基本SGD，以及具有自适应学习率的Adam优化器，再到拟二阶的L-BFGS和K-Fac优化器，最后介绍具有实际应用价值的二阶优化器AdaHessian，其可在大型推荐系统场景下比目前MLPerf打榜的AdaGrad优化器在一些数据集上拥有更快的收敛速度和更高的测试精度。

引言

目前大部分的深度学习方法都依赖于目标损失函数的构建。例如在图像分类任务中，一个损失函数通常包含神经网络输出与真实标签，它衡量了神经网络输出与真实标签之间的差距。因此，当神经网络与目标而神经网络的输出相似时，损失函数则接近于零。然而，未经过训练的神经网络输出与真实标签往往差距巨大。神经网络的输出可以通过调整网络参数来改变，因此我们需要一套能调整网络参数使得损失函数最小化的方案，而这套流程则通常称之为网络训练。如何实施神经网络的训练则是由优化器算法来描述。

优化器算法的选择是受到不同任务、不同神经网络以及不同的软硬件架构的影响。一个合适的优化器算法可以在相对短的时间内收敛到一个数值接近于零的损失函数。然而在现实中，优化过程中往往会遇到许多难题。首先，某些优化器对应的一些问题无法让损失函数曲线下降甚至导致发散，原因可能是问题结构本身并不适合一些特定的优化器算法，或者是优化器超参数例如学习率设置过高等等。其次，优化的代价过高，在大规模模型上无法实现，需要使用稀疏和分布式算法来降低内存需求。又或者优化器无法在短时间内收敛，需要使用更复杂的优化器算法来加速收敛，然而通常这些方法会额外增加计算负担，因此需要一些特殊技巧来保证优化器算法在精度和效率中得到平衡。

神经网络训练几乎是深度学习任务中最为关键一步，因为尚未训练好的网络几乎无法做出有效预测（如同抛硬币般随机的预测）。因此，随着深度学习方法的高速发展，业界与学术界也加大了对注优化器算法的关注力度，诞生了越来越多的改进算法。

梯度下降法（一阶）

在深度学习领域中，由于网络规模庞大，可调参数往往是百万、千万、上亿甚至更多。因此许多深度学习任务仍会倾向使用简洁高效计算量小的优化算法。最经典的莫过于随机梯度下降（SGD）。SGD的主要思想比较直观：不断地往梯度最陡的方向前进，持续迭代直至最终走到谷底。我们可以用数学公式来表示梯度的迭代算法

这里  指的是训练到当前迭代步数的网络参数值，  是控制前进速度的学习率，  是网络输入以及真实标签，  是网络损失函数并且也是我们想要最小化的量，而  则是损失函数对应网络参数  的梯度（深度学习任务中使用的数据集往往是非常庞大的，因此在实践中会需要将所有数据分批使用来优化网络参数，因篇幅有限，这里我们将不展开描述）。让人惊讶的是，这个简单的优化器经过简单的改良（学习率策略，加入动量项等等）仍然是目前十分主流的优化算法。原因是，它不仅具有简单且易于实现的优点，它还能训练出更具有泛化能力的神经网络，从而对未知的数据有着更好的预测能力（尤其是在计算机视觉领域）。然而，对于许多任务而言，SGD的效率较低，它需要更多的收敛步数。尤其是损失函数曲面中，相对平缓的最小值附近那些梯度接近为零的区域，SGD的收敛速度异常缓慢，当然，这可以通过增加动量项来提速，但这会造成过调而导致持续性的抖动现象而减缓收敛速度。另外，一个好的学习率以及策略也需要大量的实验来决定，不然则很有可能无法达到好的收敛结果。为了克服相关的问题，目前已有许多具有自适应能力的优化器，通过自动调整学习率来达到更好的收敛效果，例如AdaDelta，AdaGrad，RMSprop，Adam等。另外也有更进阶的优化器，例如Nadam，Radam，Lars，Lamb，NovoGrad[1]等。

牛顿优化法（二阶）

大部分的一阶优化器其实多多少少都是基于SGD上在不同方向进行改进。然而，如果想要在本质上进行收敛加速，我们则需要用到二阶导数信息来进行网络参数优化迭代。首先我们对损失函数使用泰勒级数进行函数展开

泰勒级数展开其实是对原有的函数在一个小区域使用多项式和导数进行估算，使得我们可以获得网络参数在每一步的迭代关系。这里，我们需要保证  或者每步迭代大小的值是低于1的，否则会导致泰勒级数发散而无法提供有价值的信息。通常如果我们只使用了一阶项,  的值通常需要接近为零，才能进行收敛。不过如果我们加入更高阶项时，更大的迭代步数大小也可被使用。不过代价可能是我们没有足够的计算资源来储存和计算高阶信息。因为每一阶展开项里导数的维度都比之前多出一维。例如，一阶，二阶，三阶，四阶导数对应的维度为[n]，[n，n]，[n，n，n]，[n，n，n，n]，使用Kronecker张量乘积的维度也可表达为  ，  ，  ，  ，这里n则是网络参数的数量。因此我们可以看出每一阶带来的额外计算代价都是无比巨大的。不过同样的，每一阶展开所获得的收敛加速也是非常可观的，能极大提高收敛速度。

为了找到损失函数最低值，我们需要对以上泰勒展开进行求导

这里我们设置损失函数梯度在下一步的参数值为零，原因是，在最低值时损失函数的梯度为零。这样，我们就可以找到用来更新参数的公式

这里的  是含有二阶导数信息的Hessian矩阵，我们也可以理解Hessian表达了关于损失函数的曲率。这个公式也被称为牛顿法，与SGD相比较，我们可以想象它使用了曲率信息来作为一个自适应的学习率，在越弯曲的地方走得越慢，而在完全直线的地方我们可以一步走到最底部。因此牛顿法作为一个二阶优化算法在理论上拥有着比SGD类优化器更快的收敛速度。

然而，牛顿法因为需要储存Hessian矩阵并且对其求逆的计算在深度学习任务中难以实现。例如，对于一个拥有一百万网络参数的模型而言，我们需要储存一个一兆大小的Hessian矩阵。在32位精度下，这会需要(1*10^12)*(32/8)/(10^9)=4000GB的内存！在实际应用中，例如推荐系统中一个embedding table都轻易包含了50亿以上的参数，这使得该Hessian矩阵的储存都难以被目前任何设备所储存。除此之外，我们还需要在每一步迭代中对其求逆并进行乘法操作，这是目前几乎很难高效完成的计算任务。因此目前二阶优化器算法都会在Hessian矩阵上做出一些假设或者其他技巧来进行近似估算（也称为拟牛顿法），从而大幅度地减少储存与计算量。

此外，我们还需要保证Hessian矩阵是正定的（特征值都为正数），否则牛顿法很有可能会带领参数值进入损失函数的鞍点而不是最小值。一个简单的方法则是在Hessian矩阵的对角元素加上一些量  从而使Hessian矩阵正定。我们也可以在做近似估算Hessian矩阵时构造一个接近于正定矩阵的方法，例如将要介绍的高斯-牛顿法和自然梯度法（natural gradient method）。

拟牛顿优化法（拟二阶）

因为二阶导数信息在理论上突出的收敛速度，不少二阶优化器算法已早在许多年前被不同研究者提出。大部分方法都是使用一阶梯度信息来对Hessian矩阵做近似估算。当损失函数的形式为均方误差（MSE）时（  ），我们可以使用高斯-牛顿法用根函数的平方梯度来近似Hessian矩阵（  ）。当我们将损失函数表达为对数期待值的形式时（  ），我们可以用自然梯度法使用Fisher's information matrix (FIM)来对Hessian矩阵做出近似（  ）。近期提出的K-FAC算法[2]则是利用网络结构对FIM进行分割近似。首先K-FAC法假设每一次网络参数相互不关联，使得FIM成为块矩阵，降低FIM的元素。此外，K-FAC还做出了每一份FIM矩阵块的每层网络的输出前激活梯度与上一层输出不相关的假设，因此可将其分解成两个更小的矩阵，并利用Kronecker乘积的一些特性，加大缩小需要求逆的矩阵大小。K-FAC方法除了可以在MLP上使用，还能应用在CNN[3]，RNN[4]以及其他后续提出的网络结构中。

在使用拟牛顿法进行优化时，估算出Hessian矩阵只是第一步。我们还需要算出Hessian的拟矩阵，或者至少能算出它与梯度的积  。因此也有算法直接对Hessian逆矩阵进行估算。其中最出名的莫过于BFGS算法，它使用了类似秩一校正法（rank 1 update）的方法来进行对Hessian逆矩阵（  ）的迭代。不过在大规模问题上，BFGS同样会遇到内存不够的问题，因此产生了对内存需求进行限制的L-BFGS算法，使用最近m次迭代来更新  向量。举个之前的例子，在有一百万参数的模型中，BFGS仍会需要  大小的矩阵，而L-BFGS算法只需要  大小的矩阵，这里m一般取值与10到100之间，因此大大降低内存空间。此外，还有能对参数做出约束的L-BFGS-B方法，在限制参数的数值范围的同时还能减小参数的搜索空间，进而提高收敛效率。该优化器对于神经网络在物理仿真（PINN）的任务中表现尤为出色，在速度上大大超越SGD，Adam等传统一阶算法。

AdaHessian算法（二阶）

虽然之前介绍的拟二阶优化器算法拥有着卓越的收敛效果，然而他们都在大数据的场景下表现不佳。除了计算代价过高以外，另外的一个最主要的原因在于使用mini-batch来进行参数更新对于这些方法都会造成性能损失。这里的原因是在对Hessian矩阵的近似估算缺乏鲁棒性，容易被噪声影响。

AdaHessian优化器[5]与之前提到的拟二阶优化器的不同之处在于：1）它真正意义上使用了二阶导数信息而没有做近似估算，因此拥有二阶算法的理论收敛加速。2）它只使用了Hessian矩阵的对角元素，使得二阶导数信息的维度与损失函数梯度相同，并且因此无需做出逆矩阵的计算。3）使用了Hutchinson方法和空间移动平均法来减少噪声对Hessian矩阵的影响，增加算法的鲁棒性。4）很巧妙地利用了自动微分机制来获取二阶导数信息，无需显式计算出Hessian矩阵，大大减少计算量。与一阶的梯度下降法相比，AdaHessian可以在更少步数中收敛至损失函数最小值（图1）。

图1：梯度下降法与AdaHessian在简单的  例子中的效果。梯度下降法在靠近最小值附近时因梯度接近于零的原因，需要更多迭代步数来收敛。[5]

AdaHessian的核心算法主要在于使用了Hutchinson方法对Hessian矩阵对角元素进行估算

这里  是一个随机向量拥有与参数向量一致的维度（  ），它取值与Rademacher分布（只取值与-1或1，且取其中任意一值的概率都为50%）。且拥有两个主要特性：  和  。因此，以上操作可以简单总结为  的期待值为对角线上的  项，而在非对角线上的  项的期望值则逼近于零。整个过程呈现在图2中。

图2：使用Hutchinson方法来计算Hessian矩阵的对角元素。[5]

这么计算Hessian矩阵的对角元素的好处是我们可以直接在自动微分的框架下对  进行自动微分而不需要额外计算Hessian矩阵

在得到了Hessian对角向量后，AdaHessian会对其进行移动平均(SMA)的操作来降低噪声的影响

这里  是移动平均的窗口大小。例如在卷积网络中，每一层卷积核的Hessian对角向量的大小都会有所不同，因此我们也可以将  设为一个卷积核的大小（图3）。不过，  的选值似乎在一些常见任务中对AdaHessian的整体预测效果影响不大[5]。

图3：基于卷积核大小的平均窗口选择对应Hessian对角向量的移动平均操作。[5]

我们总结下AdaHessian的算法，当得到平均过的Hessian对角向量  后，我们就可以进行参数更新了。具体操作与Adam算法类似，AdaHessian的一阶动量项  和二阶动量项  分别由以下公式进行更新

这里  分别为一阶和二阶动量系数，  为损失函数的梯度,  为可调参数，可以被默认为  。AdaHessian的参数更新公式为

 这里的  是可变化的学习率，可根据具体学习率策略来调整，例如cosine annealing。从图4中我们可以看到AdaHessian与目前主流深度学习优化器的算法对比。

图4：AdaHessian对比目前主流的优化器算法。[5]

作为一个真正意义的二阶优化器，AdaHessian每一步所使用的时间也仅仅比SGD和Adam慢两倍左右（图5）。因此与许多传统的二阶优化器不同,它可以很好地被应用于目前主流地深度学习任务中，例如视觉，语音识别，翻译，推荐系统等。

图5：AdaHessian于SGD和Adam在视觉任务中的计算效率对比

在推荐系统任务中，相对于目前主流的AdaGrad优化器，AdaHessian能意外地提高测试精度（图6）。虽然提高比例貌似不大，但是在推荐系统领域里，即使很小的精度提升都是极为困难的，并且因为用户基数巨大，因此小精度提升也往往会带来非常可观的收益[6]。

图6：AdaHessian与AdaGrad在Criteo数据集对于DLRM网络的训练与测试精度收敛图。[5]

总结

目前大部分的深度学习任务仍然首选一阶优化器，因为这些一阶方法的具有简单易用的特点。然而，由于Hessian矩阵能够提供更多的收敛信息，且随着硬件和软件的升级加速，我们相信越来越多能应用在深度学习任务中的二阶优化器算法将被挖掘出来。除了优化器算法本身，我们还可以在其他方面加速优化效果。例如，调用学习策略，或者在网络中加入正则化（权重衰减，  归一化，shake-shake正则化等）。近期诞生的sharpness-aware优化器(SAM)[7]则提出了一种正则化方式，在原有损失函数中加入其最小值宽度信息，来帮助优化器找到更宽阔平坦的损失函数最小值。这使得SAM训练过的神经网络拥有更好的泛化能力。此外，还有在输入编码器中使用哈希法以及一些其他工程技巧，能将原先需要几天才能训练好的辐射神经场（NeRF）降到5秒钟的训练[8]。

随着优化器的快速发展，相信未来神经网络的训练将更为便捷，并且训练出来的网络也将拥有更好的预测能力。

参考文献

[1] Austin Derrow-Pinion, Jennifer She, David Wong, et al. ETA Predictionwith Graph Neural Networks in Google Maps. 2021

[1] B. Ginsburg et al., “Stochastic Gradient Methods with Layer-wise Adaptive Moments for Training of Deep Networks,” May 2019.

[2] J. Martens and R. Grosse, “Optimizing Neural Networks with Kronecker-factored Approximate Curvature,” Mar. 2015.

[3] R. Grosse and J. Martens, “A Kronecker-factored approximate Fisher matrix for convolution layers,” Feb. 2016.

[4] J. Martens, J. Ba, and M. Johnson, “Kronecker-factored curvature approximations for recurrent neural networks,” 2018.

[5] Z. Yao, A. Gholami, S. Shen, M. Mustafa, K. Keutzer, and M. W. Mahoney, “ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning,” Jun. 2020.

[6] R. Wang, B. Fu, G. Fu, and M. Wang, “Deep & Cross Network for Ad Click Predictions,” Aug. 2017.

[7] P. Foret, A. Kleiner, H. Mobahi, and B. Neyshabur, “Sharpness-Aware Minimization for Efficiently Improving Generalization,” Oct. 2020.

[8] T. Müller, A. Evans, C. Schied, and A. Keller, “Instant Neural Graphics Primitives with a Multiresolution Hash Encoding,” Jan. 2022, doi: 10.1145/3528223.3530127.

 往期推荐

关于壁仞科技研究院

壁仞科技研究院作为壁仞科技的前沿研究部门，旨在研究新型智能计算系统的关键技术，重点关注新型架构，先进编译技术和设计方法学，并将逐渐拓展研究方向，探索未来智能系统的各种可能。壁仞科技研究院秉持开放的原则，将积极投入各类产学研合作并参与开源社区的建设，为相关领域的技术进步做出自己的贡献。

扫码关注我们